# -*- coding: utf-8 -*-
import re
from typing_extensions import Self
import numpy as np
import copy
from numpy.linalg import norm
from node import Node
from ElementBase import elementBase
from tools import *
from typing import List,Callable
from Loads import *
import scipy as sc

class C3D8(elementBase):
    """
    八节点六面体单元,完全积分,B-Bar技术修正体积自锁,节点自由度ux,uy,uz,节点排序和abaqus的C3d8一致;
    应力分量排序为 {11,22,33,12,23,13} ,和abaqus/explicit一致
    
    单元定义需: 节点坐标,模量E,泊松比nu,密度rho,热膨胀系数alpha
    
    支持的单元载荷: 体积力(N/m^3);压力载荷(N/m^2);表面力载荷(N/m^2)
    
    支持集中节点力加载:F1,F2,F3
    
    节点力向量,节点位移向量使用1维数组
    """
    VTK_TYPE=12

    def __init__(self, label:int,
                node1:Node,node2:Node,node3:Node,node4:Node,
                node5:Node,node6:Node,node7:Node,node8:Node,
                E:float=0.0,nu:float=0.0,rho:float=0.0,alpha:float=0.0):
        super().__init__(label=label,nodeDofNUM=3)
        self.node1,self.node2,self.node3,self.node4=node1,node2,node3,node4
        self.node5,self.node6,self.node7,self.node8=node5,node6,node7,node8
        self.E=E
        self.nu=nu
        self.rho=rho
        self.TheramlExpandFactor=alpha
        self.init_vars()
    def init_vars(self):
        """
        初始化类的内部变量,高级积分点/权重,单元面定义
        """
        self.gauss_num=2
        # 初始化单元面上的2x2高斯积分点,用于面力的等效节点力计算
        p,_=sc.special.roots_legendre(self.gauss_num)
        self.face_gauss_p={1:np.array([[i,j,-1] for i in p for j in p]),
                        2:np.array([[i,j,1] for i in p for j in p]),
                        3:np.array([[i,-1,j] for i in p for j in p]),
                        4:np.array([[1,i,j] for i in p for j in p]),
                        5:np.array([[i,1,j] for i in p for j in p]),
                        6:np.array([[-1,i,j] for i in p for j in p])}
        self.face_gauss_w={1:np.ones(4),
                        2:np.ones(4),
                        3:np.ones(4),
                        4:np.ones(4),
                        5:np.ones(4),
                        6:np.ones(4)}
        # 单元面的编号,按照右手定则会得到向内的法向
        # Face 1 1 – 2 – 3 – 4 face
        # Face 2 5 – 8 – 7 – 6 face
        # Face 3 1 – 5 – 6 – 2 face
        # Face 4 2 – 6 – 7 – 3 face
        # Face 5 3 – 7 – 8 – 4 face
        # Face 6 4 – 8 – 5 – 1 face
        self.faces_nodeseq_dict={1:[1,2,3,4],
                            2:[5,8,7,6],
                            3:[1,5,6,2],
                            4:[2,6,7,3],
                            5:[3,7,8,4],
                            6:[4,8,5,1]}
        self.faceId_dict={1:[self.node1,self.node2,self.node3,self.node4],
                        2:[self.node5,self.node8,self.node7,self.node6],
                        3:[self.node1,self.node5,self.node6,self.node2],
                        4:[self.node2,self.node6,self.node7,self.node3],
                        5:[self.node3,self.node7,self.node8,self.node4],
                        6:[self.node4,self.node8,self.node5,self.node1]}
        # 等参元的节点坐标(r,s,t)
        self.rst_coords=np.array([[-1,-1,-1],
                            [1,-1,-1],
                            [1,1,-1],
                            [-1,1,-1],
                            [-1,-1,1],
                            [1,-1,1],
                            [1,1,1],
                            [-1,1,1]],dtype=float)
        # 计算三个方向上的高斯积分点和积分权重
        rs,wrs=sc.special.roots_legendre(self.gauss_num)
        ss,wss=sc.special.roots_legendre(self.gauss_num)
        ts,wts=sc.special.roots_legendre(self.gauss_num)
        self.gauss_p=np.array([[r,s,t] for r in rs for s in ss for t in ts])
        self.gauss_w=np.array([wr*ws*wt for wr in wrs for ws in wss for wt in wts])
        if self.gauss_p.shape!=(8,3):
            raise ValueError
        if self.gauss_w.shape[0] !=8:
            raise ValueError
    @property
    def D(self)->np.ndarray:
        """
        计算单元各项同性弹性矩阵D
        
        returns:
            ndarray(6,6): 弹性矩阵D(sigama=D*epsilon)
        """
        E,nu=self.E,self.nu
        d=np.zeros((6,6))
        d[0,0]=1-nu
        d[1,1]=1-nu
        d[2,2]=1-nu
        d[0,1]=nu
        d[0,2]=nu
        d[1,2]=nu
        d[1,0]=nu
        d[2,0]=nu
        d[2,1]=nu
        d[3,3]=(1-2*nu)*0.5
        d[4,4]=(1-2*nu)*0.5
        d[5,5]=(1-2*nu)*0.5
        return (E/((1+nu)*(1-2*nu)))*d
    @property
    def get_connectivity(self):
        """
        返回单元的节点编号列表
        
        returns:
            List[int]: 单元的节点编号列表
        """
        return [self.node1.label,self.node2.label,self.node3.label,self.node4.label,
                self.node5.label,self.node6.label,self.node7.label,self.node8.label]
    @property
    def dofNum(self):
        """
        返回单元自由度总数
        
        returns:
            int: value=24
        """
        return 24
    @property
    def NodesList(self):
        """
        返回单元的节点对象列表
        
        returns:
            List[Node]: 单元的节点对象列表
        """
        return [self.node1,self.node2,self.node3,self.node4,
                self.node5,self.node6,self.node7,self.node8]
    @property
    def ElemDofIndexs(self):
        """
        单元自由度索引列表(24,)
        
        returns:
            List[int]: 单元自由度索引列表(o-based)
        """
        return dofInds(self.node1.label,3)+dofInds(self.node2.label,3)+\
            dofInds(self.node3.label,3)+dofInds(self.node4.label,3)+\
            dofInds(self.node5.label,3)+dofInds(self.node6.label,3)+\
            dofInds(self.node7.label,3)+dofInds(self.node8.label,3)
    def Ni(self,i:int)->Callable:
        """返回第i个节点的形函数Ni(函数对象,参数为r,s,t)"""
        ri,si,ti=self.rst_coords[i]
        Ni_fun=lambda r,s,t: 0.125*(1+r*ri)*(1+s*si)*(1+t*ti)
        return Ni_fun
    def Ni_r(self,i:int)->Callable:
        """返回i节点形函数Ni对r的偏导"""
        ri,si,ti=self.rst_coords[i]
        Ni_r_fun=lambda r,s,t: 0.125*ri*(1+s*si)*(1+t*ti)
        return Ni_r_fun
    def Ni_s(self,i:int)->Callable:
        """返回i节点形函数Ni对s的偏导"""
        ri,si,ti=self.rst_coords[i]
        Ni_s_fun=lambda r,s,t: 0.125*si*(1+t*ti)*(1+r*ri)
        return Ni_s_fun
    def Ni_t(self,i:int)->Callable:
        """返回i节点形函数Ni对t的偏导"""
        ri,si,ti=self.rst_coords[i]
        Ni_t_fun=lambda r,s,t: 0.125*ti*(1+r*ri)*(1+s*si)
        return Ni_t_fun
    def N(self,r,s,t)->np.ndarray:
        """输入局部坐标,返回3X24的单元形函数矩阵"""
        N_mat=np.zeros((3,24))
        for i in range(8):
            ni_fun=self.Ni(i)
            N_mat[0,3*i]=ni_fun(r,s,t)
            N_mat[1,3*i+1]=ni_fun(r,s,t)
            N_mat[2,3*i+2]=ni_fun(r,s,t)
        return N_mat
    def J(self,r,s,t)->np.ndarray:
        """计算(r,s,t)处的J矩阵(3x3)"""
        x1,y1,z1=self.node1.coordinate
        x2,y2,z2=self.node2.coordinate
        x3,y3,z3=self.node3.coordinate
        x4,y4,z4=self.node4.coordinate
        x5,y5,z5=self.node5.coordinate
        x6,y6,z6=self.node6.coordinate
        x7,y7,z7=self.node7.coordinate
        x8,y8,z8=self.node8.coordinate
        x_arr=np.array([[x1,y1,z1],
                        [x2,y2,z2],
                        [x3,y3,z3],
                        [x4,y4,z4],
                        [x5,y5,z5],
                        [x6,y6,z6],
                        [x7,y7,z7],
                        [x8,y8,z8]],dtype=float)
        N_p=np.zeros((3,8))
        for i in range(8):
            ni_r=self.Ni_r(i)
            ni_s=self.Ni_s(i)
            ni_t=self.Ni_t(i)
            N_p[0,i]=ni_r(r,s,t)
            N_p[1,i]=ni_s(r,s,t)
            N_p[2,i]=ni_t(r,s,t)
        j=np.dot(N_p,x_arr)
        return j
    def B(self,r,s,t)->np.ndarray:
        """计算坐标(r,s,t)处的B矩阵(6x24)"""
        bs=[]
        for i in range(8):
            nr=self.Ni_r(i)
            ns=self.Ni_s(i)
            nt=self.Ni_t(i)
            b=np.array([nr(r,s,t),ns(r,s,t),nt(r,s,t)])
            nx,ny,nz=sc.linalg.solve(self.J(r,s,t),b)
            bi=np.array([[nx,0,0],
                        [0,ny,0],
                        [0,0,nz],
                        [ny,nx,0],
                        [0,nz,ny],
                        [nz,0,nx]],dtype=float)
            bs.append(bi)
        b8=np.concatenate(bs,axis=1)
        return b8
    def B_modify(self,r,s,t):
        """使用b-bar技术生成B=B_dil+B_dev"""
        # 算法来源:易木响叮当,微信公众号文章
        # step1计算B矩阵
        b_dev,b_dil=[],[]
        T=np.array([[1,0,0,0,0,0],
                [0,1,0,0,0,0],
                [0,0,1,0,0,0],
                [0,0,0,0,1,0],
                [0,0,0,0,0,1],
                [0,0,0,1,0,0]])
        for i in range(8):
            nr=self.Ni_r(i)
            ns=self.Ni_s(i)
            nt=self.Ni_t(i)
            b=np.array([nr(r,s,t),ns(r,s,t),nt(r,s,t)])
            nx,ny,nz=sc.linalg.solve(self.J(r,s,t),b)
            bi=np.array([[nx,0,0],
                        [0,ny,0],
                        [0,0,nz],
                        [ny,nx,0],
                        [0,nz,ny],
                        [nz,0,nx]],dtype=float)
            # step2转换到标准B矩阵
            bi_standard=T.dot(bi)
            # step3 B矩阵变换
            bi_dil=(1/3.0)*np.array([[nx,ny,nz],
                                    [nx,ny,nz],
                                    [nx,ny,nz],
                                    [0,0,0],
                                    [0,0,0],
                                    [0,0,0]],dtype=float)
            bi_dev=bi_standard-bi_dil

            b_dev.append(bi_dev)
            b_dil.append(bi_dil)
        b8_dev=np.concatenate(b_dev,axis=1)
        b8_dil=np.concatenate(b_dil,axis=1)
        return b8_dev,b8_dil
    def detJ(self,r,s,t)->float:
        return np.linalg.det(self.J(r,s,t))
    @property
    def Ke_bbar(self):
        """返回b-bar修正后的单元刚度矩阵"""
        ke=np.zeros((24,24))
        # 对B_dev进行全积分
        for i in range(self.gauss_w.shape[0]):
            # 积分点坐标,权重
            r,s,t=self.gauss_p[i]
            w=self.gauss_w[i]
            detJ=self.detJ(r,s,t)
            B_dev,_=self.B_modify(r,s,t)
            B_devT=B_dev.T
            ke+=w*detJ*np.dot(B_devT,np.dot(self.D,B_dev))
        # B_dil进行缩减积分
        w,r,s,t=8.0,0.0,0.0,0.0
        _,B_dil=self.B_modify(r,s,t)
        B_dilT=B_dil.T
        detJ=self.detJ(r,s,t)
        ke+=w*detJ*np.dot(B_dilT,np.dot(self.D,B_dil))
        return ke
    @property
    def Ke(self)->np.ndarray:
        """单元刚度矩阵"""
        ke=np.zeros((24,24))
        for i in range(self.gauss_w.shape[0]):
            # 积分点坐标,权重
            r,s,t=self.gauss_p[i]
            w=self.gauss_w[i]
            detJ=self.detJ(r,s,t)
            BT=self.B(r,s,t).T
            ke+=w*detJ*np.dot(BT,np.dot(self.D,self.B(r,s,t)))
        return ke
    @property
    def Me_Lumped(self)->np.ndarray:
        """计算集中质量矩阵(24,24)"""
        rho=self.rho
        el_mass=0.0
        for i in range(self.gauss_w.shape[0]):
            r,s,t=self.gauss_p[i]
            el_mass+=rho*self.gauss_w[i]*self.detJ(r,s,t)
        me=(el_mass/24.0)*np.eye(24)
        return me
    @property
    def Me_Consistent(self)->np.ndarray:
        """一致质量矩阵(24,24)"""
        me=np.zeros((24,24))
        for i in range(8):
            r,s,t=self.gauss_p[i]
            w=self.gauss_w[i]
            me+=(w*self.rho*self.detJ(r,s,t))*np.dot(self.N(r,s,t).T,self.N(r,s,t))
        return me
    def ProcessELoad2EqNodeForce(self,load)->np.ndarray:
        if isinstance(load,BodyLoad):
            return self.ProcessELoad2EqNodeForceBodyLoad(load)
        if isinstance(load,PressureLoad):
            return self.ProcessELoad2EqNodeForcePressureLoad(load)
        if isinstance(load,SurfaceLoad):
            return self.ProcessELoad2EqNodeForceSurfaceLoad(load)
        if isinstance(load,ThermalLoad):
            return self.ProcessELoad2EqNodeForceThermalLoad(load)
        return np.zeros(24)
    def ProcessELoad2EqNodeForceThermalLoad(self,load:ThermalLoad)->np.ndarray:
        """计算温度载荷(摄氏度)的等效节点力"""
        eqf=np.zeros(24)
        if not isinstance(load.Temperature,float):
            raise ValueError('暂时只支持各向同性热膨胀')
        for i in range(self.gauss_w.shape[0]):
            r,s,t=self.gauss_p[i]
            w=self.gauss_w[i]
            eqf+=(w*self.TheramlExpandFactor*load.Temperature*self.detJ(r,s,t))*(np.dot(np.dot(self.B(r,s,t).T,self.D),np.array([1,1,1,0,0,0])))
        return eqf
    def ProcessELoad2EqNodeForceBodyLoad(self, load:BodyLoad) -> np.ndarray:
        """计算体积力(N/m^3)的等效节点力"""
        b=np.zeros(3)
        eqf=np.zeros(24)
        if load.Direction=='x':
            b[0]=load.LoadValue
        elif load.Direction=='y':
            b[1]=load.LoadValue
        elif load.Direction=='z':
            b[2]=load.LoadValue
        # 遍历高斯积分点
        for i in range(self.gauss_w.shape[0]):
            r,s,t=self.gauss_p[i]
            w=self.gauss_w[i]
            eqf+=(w*self.detJ(r,s,t))*(self.N(r,s,t).dot(b))
        return eqf
    def face_normal(self,faceId:int)->np.ndarray:
        """计算单元面i(faceId)的单位外法向(3,) array"""
        n1,n2,n3,n4=self.faceId_dict[faceId]
        return -1.0*calNormal(n1=n1.coordinate,n2=n2.coordinate,n3=n3.coordinate)
    def faceN(self,face_id:int,r:int,s:int,t:int)->np.ndarray:
        """计算单元面i的形函数在(r,s,t)的数值矩阵,其他的节点自由度为0"""
        N=np.zeros((3,24))
        for i in range(8):
            if i+1 in self.faces_nodeseq_dict[face_id]:
                continue
            ni=self.Ni(i)
            N[0,3*i]=ni(r,s,t)
            N[1,3*i+1]=ni(r,s,t)
            N[2,3*i+2]=ni(r,s,t)
        return N
    def node_x_y_z(self)->Tuple[np.ndarray,np.ndarray,np.ndarray]:
        """返回单元节点的x,y,z坐标数组"""
        x1,y1,z1=self.node1.coordinate
        x2,y2,z2=self.node2.coordinate
        x3,y3,z3=self.node3.coordinate
        x4,y4,z4=self.node4.coordinate
        x5,y5,z5=self.node5.coordinate
        x6,y6,z6=self.node6.coordinate
        x7,y7,z7=self.node7.coordinate
        x8,y8,z8=self.node8.coordinate
        x=np.array([x1,x2,x3,x4,x5,x6,x7,x8])
        y=np.array([y1,y2,y3,y4,y5,y6,y7,y8])
        z=np.array([z1,z2,z3,z4,z5,z6,z7,z8])
        return x,y,z
    def dr_arr(self,r,s,t)->np.ndarray:
        """返回dr向量"""
        x,y,z=self.node_x_y_z()
        nr=np.array([self.Ni_r(i)(r,s,t) for i in range(8)])
        return np.array([nr.dot(x),nr.dot(y),nr.dot(z)])
    def ds_arr(self,r,s,t)->np.ndarray:
        """返回ds向量"""
        x,y,z=self.node_x_y_z()
        ns=np.array([self.Ni_s(i)(r,s,t) for i in range(8)])
        return np.array([ns.dot(x),ns.dot(y),ns.dot(z)])
    def dt_arr(self,r,s,t)->np.ndarray:
        """返回dt向量"""
        x,y,z=self.node_x_y_z()
        nt=np.array([self.Ni_t(i)(r,s,t) for i in range(8)])
        return np.array([nt.dot(x),nt.dot(y),nt.dot(z)])
    def area_factor(self,face_id,r,s,t):
        """根据不同的单元面,生成(r,s,t)处的面积换元因子"""
        if face_id==1 or face_id==2:
            # t=+1, -1
            return np.linalg.norm(np.cross(self.dr_arr(r,s,t),self.ds_arr(r,s,t)))
        if face_id==3 or face_id==5:
            # s=+1,-1
            return np.linalg.norm(np.cross(self.dt_arr(r,s,t),self.dr_arr(r,s,t)))
        if face_id ==4 or face_id==6:
            # r=+1,-1
            return np.linalg.norm(np.cross(self.ds_arr(r,s,t),self.dt_arr(r,s,t)))
        else:
            raise ValueError("face_id error")
    def ProcessELoad2EqNodeForcePressureLoad(self,load)-> np.ndarray:
        """计算作用在单元面上的压力(N/m^2)的等效节点力,(24,)"""
        eqf=np.zeros(24)
        # 找出这个单元标签的位置在load.get_ElemsList中的所有位置
        positions = np.where(np.in1d(load.get_ElemsList, [self.label]))[0]
        # load定义中,可以一个单元的多个面都有压力,所以需要遍历
        for pos in positions:
            # 根据elem face id 进行计算
            face_id=int(load.FaceIdList[pos])
            p_arr=self.face_normal(face_id)*load.PressureValue
            # 遍历高斯积分点,使用2x2高斯积分点
            for i in range(4):
                # 根据不同的单元面生成积分点局部坐标
                r,s,t=self.face_gauss_p[face_id][i]
                w=self.face_gauss_w[face_id][i]
                face_N_matrix=self.faceN(face_id,r,s,t)
                eqf+=(w*self.area_factor(face_id,r,s,t))*(np.dot(face_N_matrix.T,p_arr))
        return eqf
    def ProcessELoad2EqNodeForceSurfaceLoad(self,load:SurfaceLoad)->np.ndarray:
        """计算面力的等效节点力(24,)"""
        eqf=np.zeros(24)
        # 找出这个单元标签的位置在load.get_ElemsList中的所有位置
        positions = np.where(np.in1d(load.get_ElemsList, [self.label]))[0]
        # load定义中,可以一个单元的多个面都有压力,所以需要遍历
        for pos in positions:
            # 根据elem face id 进行计算
            face_id=int(load.FaceIdList[pos])
            p_arr=load.DirectionVector*load.LoadValue
            # 遍历高斯积分点,使用2x2高斯积分点
            for i in range(4):
                # 根据不同的单元面生成积分点局部坐标
                r,s,t=self.face_gauss_p[face_id][i]
                w=self.face_gauss_w[face_id][i]
                face_N_matrix=self.faceN(face_id,r,s,t)
                eqf+=(w*self.area_factor(face_id,r,s,t))*(np.dot(face_N_matrix.T,p_arr))
        return eqf
    def calStrain(self,r,s,t,ue:np.ndarray)->np.ndarray:
        """get strain vector(6,), 分量11 22 33 12 23 13"""
        return self.B(r,s,t).dot(ue)
    def calStress(self,r,s,t,ue):
        """get stress vector(6,) , 应力分量11 22 33 12 23 13"""
        return self.D.dot(self.B(r,s,t)).dot(ue)
    def cal_stress_strain(self,ue:np.ndarray,s_or_e:str="stress")->np.ndarray:
        """计算积分点处stress/strain,再使用插值法外推到节点应力.返回值是一个(6,8)的数组,应力分量11 22 33 12 23 13"""
        if ue.ndim!=1:
            ue=np.reshape(ue,(-1))
        if ue.shape[0]!=self.dofNum:
            raise ValueError('dof num error')
        el_stress_at8_gauss=np.zeros((8,6))
        # 6个应力分量在8个节点的值
        el_stress_at8_node=np.zeros((6,8))
        gauss_N=np.zeros((8,8))
        # 遍历高斯点,计算应力
        for i in range(8):
            r,s,t=self.gauss_p[i]
            if s_or_e=="stress":
                el_stress_at8_gauss[i]=self.calStress(r,s,t,ue)
            if s_or_e=="strain":
                el_stress_at8_gauss[i]=self.calStrain(r,s,t,ue)
            # 填充应力插值矩阵
            for j in range(8):
                gauss_N[i,j]=self.Ni(j)(r,s,t)
        # 遍历应力分量,外推到节点
        for i in range(6):
            # 外推公式
            el_stress_at8_node[i]=sc.linalg.solve(gauss_N,el_stress_at8_gauss[:,i])
        return el_stress_at8_node

